"""
@Author: Bright
@File: callback_process.py
@Time: 2025/9/30
@Desc: 【调用过程中】回调
"""
import os

from dotenv import load_dotenv
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI, AzureChatOpenAI


class MyCustomHandler(BaseCallbackHandler):
    def on_llm_new_token(self, token: str, **kwargs) -> None:
        """做流式输出的过程中，打印每次的token"""
        print(f"My custom handler, token: {token}")


prompt = ChatPromptTemplate.from_template("给我讲个关于{animal}的笑话，限制20个字")
# 我们将自定义处理程序作为毁掉参数的列表传入，这里需要启动Stream为True时，才能出发on_llm_new_token回调。
# model = ChatOpenAI(
#     model="gpt-4", streaming=True, callbacks=[MyCustomHandler()],
# )
load_dotenv()
llm_model = AzureChatOpenAI(
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
    api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
    temperature=0.7,
    streaming=True,
    callbacks=[MyCustomHandler()]  # 我们将自定义处理程序作为毁掉参数的李彪传入
)

chain = prompt | llm_model
response = chain.invoke({"animal": "猫"})
print(response.content)
